import torch
from flask_cors import CORS
from flask_socketio import SocketIO, emit
from flask_talisman import Talisman

from trace import Trace

from flask import Flask, request, jsonify, send_from_directory, render_template, abort
import requests
import threading
import time
import trac_config
from sql_options import *
import urllib3
from utils.options import args_parser, get_config
from enum import Enum

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

app = Flask(__name__, static_folder="./dist/static/",
            template_folder="./dist")
cors_origins = os.getenv('CORS_ORIGIN', '10.3.28.134')
# CORS(app, resources={r"/*": {"origins": cors_origins}})  # 允许所有域访问所有路由
CORS(app, resources={r"/*": {"origins": "*"}})  # 允许所有域访问所有路由

csp = {
    'default-src': ["'self'"],
    'script-src': ["'self'", "https://buttons.github.io/", "https://cdnjs.cloudflare.com",
                   "https://kit.fontawesome.com", "'nonce-randomNonceValue'"],
    'style-src': ["'self'", "https://fonts.googleapis.com", "https://cdnjs.cloudflare.com", "'nonce-randomNonceValue'"],
    'img-src': ["'self'", "data:"],
    'font-src': ["'self'", "https://fonts.gstatic.com", "https://cdnjs.cloudflare.com", "https://ka-f.fontawesome.com",
                 "data:"],
    # 'frame-ancestors': ["'none'"],
    'connect-src': ["'self'", "https://ka-f.fontawesome.com"],
    'form-action': ["'self'"],
    'object-src': ["'none'"],
    'base-uri': ["'self'"],
}

# talisman = Talisman(app, content_security_policy=csp, force_https=False)
# socketio = SocketIO(app, cors_allowed_origins="*")

trace_server = None
num_clients = None
args = get_config()

class ServerStep(Enum):
    NONE = "step_none"  # 未开始
    INIT_DONE = "step_init_done"  # 初始化完成
    TRACE_DONE = "step_trace_done"  # 追踪完成
    RETURN_DONE = "step_return_done"  # 返回完成

    # 添加一个方法，返回对应的中文描述
    def get_description(self):
        descriptions = {
            ServerStep.NONE: "未开始",
            ServerStep.INIT_DONE: "初始化完成",
            ServerStep.TRACE_DONE: "追踪完成",
            ServerStep.RETURN_DONE: "返回完成",
        }
        return descriptions.get(self, "未知步骤")

def get_default_data():
    return {
        "agg_server": {
            "name": "服务器",
            "client_id": "Server",
            "status": False
        },
        "client_map": {},
        "client_map1": {},
        "step_status": ServerStep.NONE.value
    }

cache_data = get_default_data()
global conn
conn = create_connection()
if conn:
    create_trace_data_table(conn)
conn.close()


@app.route('/initial', methods=['POST'])
def initial():
    # 连接数据库
    conn = create_connection()
    global trace_server
    global num_clients
    # 更新多个值
    args.update({
        'device': torch.device('cuda:0' if torch.cuda.is_available() and 0 != -1 else 'cpu'),
        'method': 'VeriTrac',
        'global_ep': 1
    })
    # args = args_parser()
    # args.device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() and args.gpu != -1 else 'cpu')
    # args.method = 'VeriTrac'
    # args.global_ep = 1
    num_clients = count_clients(conn)
    conn.close()
    print("num_clients:  ", num_clients)
    if not num_clients:
        cache_data["step_status"] = ServerStep.NONE.value
        trace_server = None
        return jsonify({'status': 'error'}), 200
    trace_server = Trace(args)
    trace_server.func_init(args, num_clients)
    return jsonify({'status': 'success'}), 200


@app.route('/get_client_params', methods=['POST'])
def get_client_params():
    # 连接数据库
    conn = create_connection()
    client_id = request.headers.get('Client-ID')
    print("client_id", client_id)
    client_row_id = get_id(conn, client_id)
    print("client_row_id", client_row_id)
    conn.close()
    params = trace_server.get_client_params()
    params.update({
        'client_num': client_row_id,
        'num_clients': num_clients
    })
    print("client ready")
    return jsonify({'status': 'success', 'params': params}), 200


@app.route('/get_server_params', methods=['POST'])
def get_server_params():
    params = trace_server.get_server_params()
    params.update({
        'num_clients': num_clients
    })
    print("server ready")
    return jsonify({'status': 'success', 'params': params}), 200


# @app.route('/setparams', methods=['POST'])
# def setparams():
#     data = request.json
#     trace_server.set_params(data)
#     context = trace_server.get_agg_context()
#     return jsonify({'status': 'success','context': context}), 200
#
# @app.route('/get_client_context', methods=['POST'])
# def get_client_data():
#     context = trace_server.get_client_context()
#
#     return jsonify({'context': context}), 200

@app.route('/trace', methods=['POST'])
def trace():
    # 连接数据库
    conn = create_connection()
    print('trace_server.num_clients', trace_server.num_clients)
    client_id = request.headers.get('Client-ID')
    if client_id_exists(conn, client_id):
        # socketio.emit('trace_request', {'client_id': client_id})
        if cache_data["client_map"]:
            client_cache = cache_data["client_map"].get(client_id)
            if client_cache:
                client_cache["status"] = True
        trace_server.num_trace += 1
        print('trace_server.num_trace', trace_server.num_trace)
    else:
        print("No client_id found and client_id is None")
    conn.close()
    return jsonify({'status': 'success'}), 200


@app.route('/start_trace', methods=['POST'])
def start_trace():
    # 连接数据库
    conn = create_connection()
    response = requests.get(f'{trac_config.SERVER_AGG_HOST}/post_clients_data',
                            verify=False)
    clients = response.json()
    clients = clients.get('all_clients')
    for client in clients:
        insert_or_update_trace_data(conn, client['client_id'], client_ip=client['client_ip'],
                                    client_sign1=client['client_sign1'],
                                    client_sign2=client['client_sign2'], client_BFV=client['client_BFV'],
                                    pk_sign_1=client['pk_sign_1'], pk_sign_2=client['pk_sign_2'])

    global unhonest_participants
    unhonest_participants = []
    unhonest_parts = trace_server.start_trace(conn)

    client_cache = cache_data.get("client_map1")
    client_list = []
    if client_cache:
        client_list = list(client_cache.values())
    for char in unhonest_parts:
        if char.isdigit():
            unhonest_participants.append(int(char))
            if client_list and int(char) < len(client_list):
                client_list[int(char)]["status"] = True
        else:
            unhonest_participants.append(-1)
            agg_server_cache = cache_data.get("agg_server")
            if agg_server_cache:
                agg_server_cache["status"] = True
    conn.close()
    print(type(unhonest_participants[0]))
    print(unhonest_participants)
    cache_data["step_status"] = ServerStep.TRACE_DONE.value
    return jsonify({'unhonest_participants': unhonest_participants}), 200


@app.route('/return_result', methods=['POST'])
def return_result():
    # 连接数据库
    conn = create_connection()
    data = request.json
    clients = data.get('request_trace_ids')
    result_array = []
    for participant in unhonest_participants:
        if participant == -1:
            result_array.append("Server")
        elif participant > -1:
            # id = get_client_id(conn,participant+1)
            client_info = get_client_info(conn, participant + 1)

            if id:
                # result_array.append(id[0])
                result_array.append(client_info)

            else:
                result_array.append(None)

    for request_client_id in clients:
        clientdatas = get_basedata_trace_clients_data(conn, request_client_id)
        for client in clientdatas:
            client_ip = client[0]
            # client_port = client[1]

        print(f"Client ID: {request_client_id}, IP: {client_ip}")
        headers = {'Content-Type': 'application/json', 'Client-ID': request_client_id}
        data = {
            'unhonest_participant': result_array,
        }
        response = requests.post(
            f'{client_ip}/trace_result',
            json=data, headers=headers, verify=False)
        print(f"Sent data to client: {response.json()}")
    conn.close()
    init_data()
    cache_data["step_status"] = ServerStep.RETURN_DONE.value
    return jsonify({'status': 'success'}), 200


def agg_server_init():
    response = requests.post(f'{trac_config.SERVER_AGG_HOST}/reset', verify=False)
    if response.status_code == 200:
        print("agg_server init successfully!")
    else:
        print("agg_server init error: {}".format(response.text))


def client_batch_logout():
    try:
        conn = create_connection()
        clients_simpledata = get_simpledata_trace_clients_data(conn)
        print("clients_simpledata: ", clients_simpledata)
        for item in clients_simpledata:
            client_logout(item["client_ip"], item["client_id"])
    except Exception as e:
        print("client_batch_logout error: {}".format(e))
    finally:
        conn.close()


def client_logout(client_ip, client_id):
    response = requests.post(f'{client_ip}/logout?client_id={client_id}', verify=False)
    if response.status_code == 200:
        print("client id [{}] , ip [{}] logged out successfully".format(client_id, client_ip))
    else:
        print("client id [{}] , ip [{}] logged out error".format(client_id, client_ip))
        print("error: {}".format(response.text))


@app.route('/delete_client', methods=['DELETE'])
def delete_client():
    client_id = request.json.get('client_id')
    print(f"Client ID: {client_id}")
    if not client_id:
        return jsonify({'error': 'client_id is required'}), 400
    deleted_rows = delete_client_data(client_id)
    print(deleted_rows)
    if deleted_rows is None:
        return jsonify({'error': 'An error occurred while deleting the data'}), 500
    elif deleted_rows > 0:
        return jsonify({'message': 'Client data deleted successfully'}), 200
    else:
        return jsonify({'error': 'Client ID not found'}), 404


@app.route('/data', methods=['POST'])
def receive_data():
    # 连接数据库
    conn = create_connection()
    data = request.json
    print(f"Received data: {data}")
    client_id = request.headers.get('Client-ID')
    client_ip = data.get('client_ip')
    # client_port = data.get('client_port')
    userid = data.get('userid')
    username = data.get('username')
    print(f"Received data from client {client_id}: {data}")
    insert_or_update_trace_data(conn, client_id=client_id, client_ip=client_ip,
                                userid=userid, username=username)
    conn.close()
    client_id = request.headers.get('Client-ID')
    print(f"Received data from client {client_id}: {data}")
    return jsonify({'status': 'success'}), 200


@app.route('/init', methods=['POST'])
def init():
    clients_simpledata, cache = init_data()
    cache_data["step_status"] = ServerStep.INIT_DONE.value
    return jsonify({'clients_simpledata': clients_simpledata})


def init_data():
    """
    初始化数据,并返回缓存和客户端列表
    return: 客户端数据, 缓存数据
    """
    global cache_data
    # 连接数据库
    conn = create_connection()
    # 前端初始化
    set_id(conn)
    clients_simpledata = get_simpledata_trace_clients_data(conn)
    conn.close()
    # "id, client_id, userid, username"
    client_map = {}
    client_map1 = {}
    cache_data = get_default_data()
    for item in clients_simpledata:
        def gen():
            return {
            "name": item["username"],
            "client_id": item["client_id"],
            "status": False,
            "userid": item["userid"]
        }
        client_map[item["client_id"]] = gen()
        client_map1[item["client_id"]] = gen()
    cache_data["client_map"] = client_map
    cache_data["client_map1"] = client_map1
    return clients_simpledata, cache_data


@app.route('/ping', methods=['GET'])
def ping():
    client_id = request.headers.get('Client-ID')
    return jsonify({'status': 'alive', 'client_id': client_id}), 200


@app.route('/reset', methods=['POST'])
def reset_server():
    client_batch_logout()
    agg_server_init()
    global trace_server, num_clients, cache_data
    trace_server = None
    num_clients = None
    cache_data = get_default_data()
    return jsonify({'status': 'success'}), 200


def run_flask():
    # app.run(host='0.0.0.0', port=5001, ssl_context=(
    #     f'{Path(__file__).parent}/certs/server.crt',
    #     f'{Path(__file__).parent}/certs/server.key'
    # ))
    app.config['DEBUG'] = False
    app.run(host='0.0.0.0', port=5001)
    # socketio.run(app, host='0.0.0.0', port=5001, allow_unsafe_werkzeug=True)


# WebSocket 连接事件
# @socketio.on('connect')
# def handle_connect():
#     print('Client connected')
#
#
# @socketio.on('disconnect')
# def handle_disconnect():
#     print('Client disconnected')


def connect_to_agg():
    while True:
        try:
            response = requests.get(f'{trac_config.SERVER_AGG_HOST}/ping',
                                    verify=False)
            print(f"Connected to agg server: {response.json()}")
            initial()
            break
        except requests.exceptions.RequestException:
            print("Connection failed, retrying...")
            time.sleep(5)


# 路由处理 favicon.ico 请求
@app.route('/favicon.ico')
def favicon():
    return send_from_directory(os.path.join(app.root_path, 'static'), 'favicon.png', mimetype='image/png')


@app.route('/')
@app.route('/<path:path>')
def serve(path='index.html'):
    if path != '' and os.path.exists(os.path.join(app.template_folder, path)):
        # 检查路径是否指向模板文件（如 .html 文件）
        if path.endswith('.html'):
            return render_template(path)
        else:
            # 如果路径指向静态文件，则返回静态文件
            return send_from_directory(app.template_folder, path)
    else:
        return render_template('index.html')


@app.route('/trace_status', methods=['GET'])
def trace_status():
    """
    返回 server trace 状态及会话缓存
    """
    step_status = request.args.get('step_status')
    if step_status:
        return {"step_status": cache_data["step_status"]}

    client_list = cache_data.get("client_map").values() if cache_data.get("client_map") else []
    client_list1 = cache_data.get("client_map1").values() if cache_data.get("client_map1") else []
    result = [cache_data.get("agg_server")]
    result.extend(client_list1)
    data = {
        "traceClients": list(client_list),
        "traceResult": result,
        "step_status": cache_data["step_status"]
    }
    return data


# 定义禁止访问的主机列表
BLOCKED_HOSTS = ['169.254.169.254', 'aws.zaproxy.org', '100.100.100.200', 'alibaba.zaproxy.org']


@app.before_request
def block_metadata_access():
    if request.host in BLOCKED_HOSTS:
        abort(403)  # 禁止访问


if __name__ == "__main__":
    # 检查和初始化数据库
    conn = create_connection()
    if conn:
        create_trace_data_table(conn)
    conn.close()
    # flask_thread = threading.Thread(target=run_flask)
    # flask_thread.start()
    run_flask()
    # connect_thread = threading.Thread(target=initial)
    # connect_thread.start()

    # Wait for threads to complete
    # connect_thread.join()
    # flask_thread.join()
    # conn.close()

    # global args
    # args = args_parser()
    # global conn
    # conn = create_connection()
    # if conn:
    #     create_trace_data_table(conn)

    # app.config['DEBUG'] = False
    # app.run(host='0.0.0.0', port=5003)
    # socketio.run(app, host='0.0.0.0', port=5001)
    # conn.close()
